package com.wzh.test.huffman;

import java.io.UnsupportedEncodingException;
import java.util.*;

public class HuffmanDemo {
    public static void main(String[] args) throws UnsupportedEncodingException {
        //给你一个字符串 "i like like like java do you like java"
        //要求：计算出该字符串的huffman编码
        //什么是huffman编码？
        //正常编码是使用ascii码  i-->01101001  空格-->00100000  l-->01101100  k-->01101011
        //huffman编码：根据每个字符的使用频率作为权重，生成一颗huffman树，然后左侧路径为0，右侧路径为1,生成编码
        String str = "i like like like java do you like a java";
        byte[] strBytes = str.getBytes();


        Node huffmanTree = createHuffmanTree(strBytes);
        //根据huffman树，生成编码map
        HashMap<Byte,String> huffmanCodeMap = new HashMap<>();
        StringBuilder stringBuilder = new StringBuilder();

        createHuffmanCodeMap(huffmanCodeMap,huffmanTree,stringBuilder,"");

        byte[] huffmanCodeBytes = createHuffmanCodeBytes(huffmanCodeMap,strBytes);
        System.out.println(Arrays.toString(huffmanCodeBytes));

    }

    public static byte[] createHuffmanCodeBytes(HashMap<Byte,String> huffmanCodeMap,byte[] strBytes) throws UnsupportedEncodingException {
        //System.out.println(huffmanCodeMap);
        StringBuilder sbEncoder = new StringBuilder();
        for (byte strByte : strBytes) {
            sbEncoder.append(huffmanCodeMap.get(strByte));
        }
        //System.out.println(sbEncoder.toString());
        byte[] huffmanCodeBytes = convertString2Bytes(sbEncoder);

        return huffmanCodeBytes;
    }

    public static Node createHuffmanTree(byte[] strBytes){
        //1、统计出每个字符的使用频率，放到一个map中
        HashMap<Byte,Integer> charCountMap = new HashMap<>();
        for (byte strByte : strBytes) {
            if(charCountMap.containsKey(strByte)){
                //如果有这个字符
                Integer count = charCountMap.get(strByte);
                charCountMap.put(strByte,count + 1 );
            }else{
                //如果没有这个字符
                charCountMap.put(strByte,1);
            }
        }

        //定义一个PriorityQueue<Node> 存放节点
        PriorityQueue<Node> queue = new PriorityQueue<>();
        for (Map.Entry<Byte, Integer> set : charCountMap.entrySet()) {
            Node node = new Node(set.getKey(),set.getValue());
            queue.add(node);
        }

        Node huffmanTree = createHuffmanTree(queue);
        return huffmanTree;
    }
    public static  byte[] convertString2Bytes(StringBuilder stringBuilder) throws UnsupportedEncodingException {
        int len;
        if(stringBuilder.length() % 8 == 0){
            len = stringBuilder.length() / 8;
        }else{
            len = stringBuilder.length() / 8 + 1;
        }
        //创建存储压缩后的byte数组
        byte[] huffmanCodeBytes = new byte[len];
        int index = 0;//记录是第几个byte
        for(int i = 0; i < stringBuilder.length();i+=8){
            //因为是每8位对应一个byte,所以步长+8
            String strByte;
            if(i + 8 < stringBuilder.length()) {
                strByte = stringBuilder.substring(i, i + 8);
            }else{
                strByte = stringBuilder.substring(i);
            }
            huffmanCodeBytes[index] = (byte) Integer.parseInt(strByte,2);
            index++;
        }
        return huffmanCodeBytes;
    }
    public static void createHuffmanCodeMap(HashMap<Byte,String> huffmanCodeMap,Node node,StringBuilder stringBuilder,String code){
        StringBuilder stringBuilder1 = new StringBuilder(stringBuilder);
        if(!code.equals("")) {
            stringBuilder1.append(code);
        }
        if(node.mChar != null){
            huffmanCodeMap.put(node.mChar,stringBuilder1.toString());
        }
        if(node.left != null){
            createHuffmanCodeMap(huffmanCodeMap,node.left,stringBuilder1,"0");
        }
        if(node.right != null){
            createHuffmanCodeMap(huffmanCodeMap,node.right,stringBuilder1,"1");
        }
    }

    /**
     * 功能：给你一个数组[3,4,5,6,7,1,2],要求生成一颗huffman树，并打印出来
     */
    public static Node createHuffmanTree(PriorityQueue<Node> queue) {


        while (queue.size() > 1) {
            //拿出前两个最小的节点，生成一颗树
            Node first = queue.poll();
            Node second = queue.poll();
            Node tmpRootNode = new Node(null,first.weight + second.weight);
            tmpRootNode.left = first;
            tmpRootNode.right = second;
            //把新生成的树加入到queue
            queue.add(tmpRootNode);
        }

        Node rootNode = queue.poll();
        return rootNode;
    }

    /**
     * 前序遍历打印出huffman树
     *
     * @param node
     */
    public static void preOrder(Node node) {
        System.out.println(node);
        if (node.left != null) {
            preOrder(node.left);
        }
        if (node.right != null) {
            preOrder(node.right);
        }
    }

}


class Node implements Comparable<Node> {
    public Byte mChar;
    public Integer weight;
    public Node left;
    public Node right;

    public Node(Byte mChar, Integer weight) {
        this.mChar = mChar;
        this.weight = weight;
    }


    @Override
    public String toString() {
        return "Node{" +
                "mChar=" + mChar +
                ", weight=" + weight +
                '}';
    }

    @Override
    public int compareTo(Node o) {
        return this.weight - o.weight;
    }
}
